package markov;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

import everything.WorldMap;


//how many states there are?
// 3^icyPatches*2^flags*xPosition*yPosition*direction

public class MDPState implements Comparable<MDPState>{
	public static final byte NORTH = 0;
	public static final byte SOUTH = 1;
	public static final byte EAST = 2;
	public static final byte WEST = 3;
	public static final byte NO_DIRECTION = 4;
	private static final float RewardF = 10;
	private static final float CostM = -5;
	private static final float CostR = -4.999f;
	//	private static final int CostS = -10;
	public static final byte ICE = 0;
	public static final byte NO_ICE = 1;
	public static final byte UNKNOWN = 2;
	public static final byte REST = 4;
	boolean[] flags;
	byte[] icyPatches;
	int xPosition;
	int yPosition;
	byte direction;

	int boardLength;
	int boardWidth;

	public MDPState(boolean[] flags, byte[] icyPatches, int xPosition, int yPosition, byte direction, int boardLength, int boardWidth){
		this.flags = flags;
		this.icyPatches = icyPatches;
		this.xPosition = xPosition;
		this.yPosition = yPosition;
		this.direction = direction;
		this.boardLength = boardLength;
		this.boardWidth = boardWidth;
	}

	public MDPState(int id, int numFlags, int numIces, int boardLength, int boardWidth){
		flags = new boolean[numFlags];
		icyPatches = new byte[numIces];
		for (int i=0; i<flags.length; i++){
			flags[i] = (id % 2 == 1);
			id = id / 2;
		}
		for (int i=0; i<icyPatches.length; i++){
			icyPatches[i] = (byte)(id % 3);
			id = id / 3;
		}
		xPosition = id % boardLength;
		id = id / boardLength;
		yPosition = id % boardWidth;
		id = id / boardWidth;
		direction = (byte)id;

		this.boardLength = boardLength;
		this.boardWidth = boardWidth;
	}

	public Map<MDPState, Float> nextState(byte action, WorldMap map, float[][] iceProbs){
		Map<MDPState, Float> ans = new HashMap<MDPState, Float>();

		byte movingDirection = action;
		if (direction!=NO_DIRECTION){
			movingDirection = direction;
		}
		int nextX = xPosition;
		int nextY = yPosition;
		switch(movingDirection){
		case NORTH:
			if (!map.isMovementBlocked(xPosition-1, yPosition)){
				nextX--;
			}
			break;
		case SOUTH:
			if (!map.isMovementBlocked(xPosition+1, yPosition)){
				nextX++;
			}
			break;
		case EAST:
			if (!map.isMovementBlocked(xPosition, yPosition+1)){
				nextY++;
			}
			break;
		case WEST:
			if (!map.isMovementBlocked(xPosition, yPosition-1)){
				nextY--;
			}
			break;
		}


		boolean[] newFlags = Arrays.copyOf(flags, flags.length);
		byte[] newIcyPatches = Arrays.copyOf(icyPatches, icyPatches.length);
		byte newDirection = NO_DIRECTION;

		if (map.brickStatus(nextX, nextY)==WorldMap.FLAG){
			//flag reward
			int id = map.getFlagId(nextX, nextY);
			newFlags[id] = false;
		}

		if (map.isIce(nextX, nextY)){
			//the block we are reaching is a potential ice block
			int iceId = map.getIceId(nextX, nextY);
			int iceStat = icyPatches[iceId];
			if (iceStat == NO_ICE){
				MDPState m1 = new MDPState(newFlags, newIcyPatches, nextX, nextY, newDirection, boardLength, boardWidth);
				ans.put(m1, 1.0f);
			} else if (iceStat == ICE){
				byte newDirection2 = getNewDirection(nextX, nextY, movingDirection, map);
				MDPState m2 = new MDPState(newFlags, newIcyPatches, nextX, nextY, newDirection2, boardLength, boardWidth);
				ans.put(m2, 1.0f);
			} else {	// unknown
				int evidenceId = getEvidenceId(icyPatches);
				float iceProb = iceProbs[evidenceId][iceId];

				byte newDirection2 = getNewDirection(nextX, nextY, movingDirection, map);
				byte[] newIcyPatches2 = Arrays.copyOf(icyPatches, icyPatches.length);

				newIcyPatches[iceId] = NO_ICE;
				newIcyPatches2[iceId] = ICE;

				MDPState m1 = new MDPState(newFlags, newIcyPatches, nextX, nextY, newDirection, boardLength, boardWidth);
				MDPState m2 = new MDPState(Arrays.copyOf(newFlags, newFlags.length), newIcyPatches2, nextX, nextY, newDirection2, boardLength, boardWidth);
				ans.put(m1, (1-iceProb));
				ans.put(m2, iceProb);
			}
			//TODO - sensing alone
			//TODO - sensing when on ice
			//TODO for each possible ice block - if we dont have a second sensing for it - we need to add a node
		} else {
			MDPState m1 = new MDPState(newFlags, newIcyPatches, nextX, nextY, newDirection, boardLength, boardWidth);
			ans.put(m1, 1.0f);
		}
		if (action<=0){	//sensing

		}

		return ans;
	}

	private int getEvidenceId(byte[] icyPatches) {
		int ans = 0;
		int pow = 1;
		for (int i=0; i<icyPatches.length; i++){
			ans = ans + pow*icyPatches[i];
			pow = pow*3;
		}
		return ans;
	}

	public byte getNewDirection(int nextX, int nextY,byte movingDirection, WorldMap map){
		int nextNextX = nextX;
		int nextNextY = nextY;
		switch(movingDirection){
		case NORTH:
			if (!map.isMovementBlocked(xPosition-1, yPosition)){	nextNextX--;
			} break;
		case SOUTH:
			if (!map.isMovementBlocked(xPosition+1, yPosition)){ nextNextX++;
			} break;
		case EAST:
			if (!map.isMovementBlocked(xPosition, yPosition+1)){	nextNextY++;
			} break;
		case WEST:
			if (!map.isMovementBlocked(xPosition, yPosition-1)){	nextNextY--;
			} break;
		}
		if (map.isMovementBlocked(nextNextX, nextNextY)){
			return NO_DIRECTION;
		} else {
			return movingDirection;
		}
	}

	public int getId(){
		int ans = 0;
		int pow = 1;
		for (int i=0; i<flags.length; i++){
			ans = ans + (flags[i] ? 1 : 0)*pow;
			pow = pow*2;
		}
		for (int i=0; i<icyPatches.length; i++){
			ans = ans + (icyPatches[i])*pow;
			pow = pow*3;
		}
		ans = ans + xPosition*pow;
		pow = pow*boardLength;
		ans = ans + yPosition*pow;
		pow = pow*boardWidth;
		ans = ans + direction*pow;
		return ans;
	}

	public static float reward(MDPState prev, int action, MDPState next){
		float reward = 0;
		if (action==NORTH || action==SOUTH || action==EAST || action==WEST){
			reward = reward +CostM;
			//		} else if (action==SENSE){
			//			reward = rewrad + CostS;
		} else if (action == REST){
			reward = reward + CostR;
		}
		for (int i=0 ;i<prev.flags.length; i++){
			if (!next.flags[i] && prev.flags[i]){	//flag changed from 'not taken' to 'taken'
				reward = reward + RewardF;
			}
		}
		return reward;
	}

	@Override
	public boolean equals(Object o) {
		if (o instanceof MDPState){
			return getId()==((MDPState)o).getId();
		} else return false;
	}

	@Override
	public int compareTo(MDPState o) {
		long myId = getId();
		long otherId = o.getId();
		if (myId > otherId){
			return -1;
		} else if (myId < otherId){
			return 1;
		} else {
			return 0;
		}

	}

	public String toString(){
		String ans = "State " + getId() + ": [";
		for (int i=0; i<flags.length-1; i++){
			if (flags[i]){
				ans = ans + "T,";
			} else {
				ans = ans + "F,";
			}
		}
		if (flags.length>0){
			if (flags[flags.length-1]){
				ans = ans + "T] - [";
			} else {
				ans = ans + "F] - [";
			}
		} else {
			ans = ans + "]";
		}
		for (int i=0; i<icyPatches.length-1; i++){
			ans = ans + icyPatches[i]+",";
		}
		if (icyPatches.length>0){
			ans = ans + icyPatches[icyPatches.length-1]+"]";
		} else {
			ans = ans + "]";
		}
		ans = ans + " - ["+xPosition +"," +yPosition+"]";
		ans = ans + " - " + direction;
		return ans;
	}
}

